# -*- coding: utf-8 -*-
"""
Created on Tue Apr 21 14:47:14 2020

@author: Farman
"""

import numpy as np
import PIL.Image as img  # for image manipulation
import os
import sys
import joblib
import pathlib



def parse_groud_truth_table(gt_class_file):
    '''
    Parameters
    ----------
    gt_class_file : contains groud truth class information, 
                    string in format as below.
        
        1 	Water 	                824 
        2 	Trees 	                820 
        3 	Asphalt 	            816 
        4 	Self-Blocking Bricks 	808 
        5 	Bitumen 	            808 
        6 	Tiles 	               1260 
        7 	Shadows 	            476 
        8 	Meadows 	            824 
        9 	Bare Soil 	            820

    Returns
    -------
        Dictionary in format of class {index : name}.

    '''
    with open(gt_class_file) as f:
        lines = f.readlines()
        
    dict_class = {}
    
    for line in lines:
        line = line.strip()
        [index, name, *padding] = line.split(' ')
        index = index.strip()
        name  = name.strip()
        dict_class.update({index : name})
    
    dict_class.update({0 : 'unknown'})
    
    return dict_class
        


def read_groud_truth(gt_file):
    '''
    Read groud truth image data from MyPersistence persistence file. 

    Parameters
    ----------
    gt_file : string
        Groud truth file name.

    Returns
    -------
    Groud truth data in format of numpy array(2-D).

    '''
    return joblib.load(gt_file)
    


def generate_groud_truth_mask_image(gt_data, gt_class, save_path):
    '''
    Generate ground truth mask image.
    One mask image for each gound truth value.
    Mask images saved with name 'gtvalue-gtname.png' in working directory.

    Parameters
    ----------
    gt_data : 2-D numpy array contains groud truth data.
        
    gt_glass : Dictionary of ground truth index and name, see parse_groud_truth_table.
    
    save_path : string, path for save mask image.

    Returns
    -------
    None.

    '''
    save_path = pathlib.Path(save_path)
    
    try:
        os.mkdir(save_path)
    except:
        pass
    
    for index in gt_class.keys():
        gt_data = np.array(gt_data)
        mask = 255 * (gt_data == int(index))
        mask = np.array(mask, dtype='uint8')
        pic = img.fromarray(mask)
    
        name = "%02d-%s.png"%(int(index), gt_class.get(index))
        print('Saving mask :', save_path / name)
        pic.save(save_path / name)
    
    return



def process(gt_file, gt_class_file):
    '''
    Groud truth mask images generate for outer use.

    Parameters
    ----------
    gt_file : string.
        Groud truth file name.
        
    groudtruth_table_file : string.
        Name of file contains groud truth class information.
        See parse_groud_truth_table() for details.

    Returns
    -------
    None.

    '''
    gt_class = parse_groud_truth_table(gt_class_file)
    save_path = pathlib.Path(gt_file).parent
    gt_data_dict = read_groud_truth(gt_file)
    
    for k in gt_data_dict.keys():
        gt_data = gt_data_dict.get(k)
        generate_groud_truth_mask_image(gt_data, gt_class, save_path / k)
    return
    
    

#------------------------------------------------------------------------------


if __name__ == '__main__':
    print('-' * 40, '\n')
    print('    GroundTruth Mask Image Generator\n')
    print('-' * 40, '\n')
    
    #process(r'D:\Pavia\Pavia\PaviaUniversity\PaviaU_gt.joblib', 
    #        r'D:\Pavia\Pavia\PaviaUniversity\gtClass.txt')
    
    if len(sys.argv) == 1:        
        gt_file = input('Ground Truth Image Persistence file : ')
        gt_class_file = input('Ground Truth Class File : ')
            
        if len(gt_file) and len(gt_class_file):
            process(gt_file, gt_class_file)
        else:
            print('Invlid input.')

    elif len(sys.argv) == 3:
        gt_file = sys.argv[1]
        gt_class_file = sys.argv[2]
        process(gt_file, gt_class_file)
    
    else:
        print('Show me the proper parameters.\nUsage:')
        print('    python3 GroudTruthMaskGen.py ')
        print('or')
        print('    python3 GroudTruthMaskGen.py MyPersistenceFile GTClassFile\n')